{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compose Two Different Static Maze2D Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os, sys, torch\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "sys.path.append('.') ## please add the project root dir accordingly\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.backends.cudnn.allow_tf32 = True\n",
    "torch.backends.cudnn.deterministic = True\n",
    "import diffuser.utils as utils\n",
    "from diffuser.models import GaussianDiffusionPB\n",
    "from diffuser.guides.policies_compose import PolicyCompose\n",
    "from datetime import datetime\n",
    "import os.path as osp\n",
    "from tqdm import tqdm\n",
    "from diffuser.utils.comp.comp_setup import ComposedParser\n",
    "from diffuser.utils.jupyter_utils import suppress_stdout\n",
    "from tap import Tap\n",
    "import numpy as np\n",
    "np.set_printoptions(precision=3)\n",
    "\n",
    "#---------------------------------- setup ----------------------------------#\n",
    "\n",
    "## You need to download 'rm2d-comp' from OneDrive Link in README.md to launch\n",
    "config = \"config/comp/Comp_rSmaze_nw6nw3_hExt0507_engy_add.py\"\n",
    "class Parser(Tap):\n",
    "    config: str = config\n",
    "    plan_n_maze: int = 0 # placeholder\n",
    "\n",
    "from diffuser.guides.comp_plan_env import ComposedEnvPlanner\n",
    "from diffuser.guides.comp_plan_helper import ComposedDiffusionPlanner\n",
    "\n",
    "sys.argv = ['PlaceHolder.py', '--config', config]\n",
    "\n",
    "with suppress_stdout():\n",
    "    args_comp = Parser().parse_args()\n",
    "    cp = ComposedParser(args_comp.config)\n",
    "    args_train_list, args_list = cp.setup_args_list(args_comp,)\n",
    "    dplanner = ComposedDiffusionPlanner(args_train_list, args_list, cp.comp_dataset)\n",
    "\n",
    "plan_env_list = dplanner.plan_env_list\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create an Env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffuser.utils.rm2d_render import RandStaticMazeRenderer\n",
    "import imageio\n",
    "# import pb_diff_envs.environment.static.rand_maze2d_renderer as rmr; reload(rmr)\n",
    "from pb_diff_envs.environment.static.rand_maze2d_renderer import RandMaze2DRenderer\n",
    "rootdir = './visualization/'\n",
    "\n",
    "bs = 50\n",
    "env_source = 'dataset' # or 'custom'\n",
    "env_source = 'custom'\n",
    "\n",
    "if env_source == 'custom':\n",
    "    env_id = 99 # placeholder\n",
    "    wall_locations_c = [\n",
    "        np.array([[1.13 , 3.184],\n",
    "            [1.571, 1.849],\n",
    "            [2.857, 4.093],\n",
    "            [3.4  , 2.394],\n",
    "            [3.563, 1.143],\n",
    "            [4.333, 3.949],]),\n",
    "        np.array([\n",
    "            [1.069, 1.257],\n",
    "            [3.101, 2.45],\n",
    "            [2.714, 4.14 ]]),\n",
    "        ]\n",
    "    \n",
    "    ## size of the obstacles 1x1 square block\n",
    "    if len(wall_locations_c) == 2:\n",
    "        hExt_list_c = [np.array([[0.5, 0.5],]*6), np.array([[0.7, 0.7],] * len(wall_locations_c[1]) )]\n",
    "    else:\n",
    "        hExt_list_c = [np.array([[0.5, 0.5],]*6), ]\n",
    "    env = plan_env_list.create_env_by_pos(env_id, wall_locations_c, hExt_list_c)\n",
    "    \n",
    "elif env_source == 'dataset':\n",
    "    env_id = 30 # input an index\n",
    "    env = plan_env_list.create_single_env(env_id)\n",
    "\n",
    "\n",
    "n_walls_1 = len(env.wloc_list_c[0]) ## num of the smaller 1x1 blocks\n",
    "n_walls_2 = len(env.wloc_list_c[1]) ## num of the larger 1.4x1.4 blocks\n",
    "\n",
    "obs_start_np = np.array( [ [ 1.20, 0.2 ] ], dtype=np.float32 ).repeat( bs, 0 )\n",
    "target_np = np.array( [ [3.89, 4.55] ], dtype=np.float32  ).repeat( bs, 0 )\n",
    "\n",
    "\n",
    "e_wlocs = env.wall_locations\n",
    "env_hExts = env.wall_hExts # (num_walls, 2)\n",
    "print(f'[Training Time] model 1 only trained on 3 obstacles (o); model 2 # of obstacles: 6')\n",
    "print(f'[Now] # of walls: {len(e_wlocs)}')\n",
    "\n",
    "mz_renderer = RandMaze2DRenderer(num_walls_c = list(map(lambda x: len(env.wloc_list_c[x]), list(range(len(env.wloc_list_c))) )), fig_dpi=200)\n",
    "mz_renderer.up_right = 'empty'\n",
    "mz_renderer.config['no_draw_traj_line'] = True\n",
    "mz_renderer.config['no_draw_col_info'] = True\n",
    "\n",
    "img = mz_renderer.renders( np.concatenate([obs_start_np[0:1,], target_np[0:1,],]), plan_env_list.maze_size, e_wlocs, env_hExts )\n",
    "utils.plt_img(img, dpi=80)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print( env.wall_locations.shape )\n",
    "print( env.wall_locations )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "import imageio\n",
    "import diffuser.guides.comp_plan_env as dcpe; reload(dcpe)\n",
    "from diffuser.guides.comp_plan_env import ComposedEnvPlanner, ComposedRM2DReplanner\n",
    "\n",
    "plan_env_list = dplanner.plan_env_list\n",
    "rm2d_planner = ComposedEnvPlanner()\n",
    "\n",
    "policy = dplanner.policy\n",
    "policy.ddim_eta = 1.0\n",
    "print(f' policy: {policy.cg_w}')\n",
    "use_ddim = True\n",
    "repl_dn_steps = 5\n",
    "\n",
    "wloc_np = env.wall_locations ## important\n",
    "policy.comp_type = 'share-model2fix'\n",
    "# n_walls_1 is the smaller block\n",
    "policy.num_walls_c = torch.tensor([ min(6, n_walls_1+n_walls_2), n_walls_2 ])\n",
    "print(f'policy.num_walls_c: {policy.num_walls_c}')\n",
    "\n",
    "repl_config = dict(use_ddim=use_ddim, dn_steps=repl_dn_steps)\n",
    "replanner = ComposedRM2DReplanner(policy, dplanner.train_normalizer_list[0], repl_config, 'cuda')\n",
    "\n",
    "\n",
    "obs_start_arr = utils.to_torch(obs_start_np,  device='cpu')\n",
    "target_arr = utils.to_torch( target_np, device='cpu' )\n",
    "wloc_input = wloc_np[None, ].repeat(bs, axis=0) # (10, 2) -> (bs, 10, 2)\n",
    "wloc_input = wloc_input.reshape(bs, -1)\n",
    "wall_locations = utils.to_torch( wloc_input,  device='cpu' )\n",
    "\n",
    "\n",
    "use_normed_wallLoc = False\n",
    "plan_env_config = utils.Config(None,\n",
    "    env_mazelist=dplanner.plan_env_list,\n",
    "    prob_permaze=1,\n",
    "    samples_perprob=bs,\n",
    "    obs_selected_dim=(0,1),\n",
    "    horizon_pl_list=[48, 52, 56, 60, 64],\n",
    "    horizon_wtrajs=48,\n",
    "    do_replan=False,\n",
    "    replanner=replanner,\n",
    ")\n",
    "with suppress_stdout():\n",
    "    vis_modelout_path_list, samples_mulho, _ = rm2d_planner.plan_env_interact(\n",
    "        env, policy, obs_start_arr, target_arr, wall_locations, use_normed_wallLoc, plan_env_config,)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffuser.utils.jupyter_utils import get_all_suc_trajs\n",
    "\n",
    "all_candidates = []\n",
    "for trajs in samples_mulho:\n",
    "    all_candidates.extend(trajs.observations)\n",
    "## pick out successful trajectories\n",
    "print('len all_candidates:', len(all_candidates))\n",
    "env.robot.collision_eps = 1e-2\n",
    "suc_trajs, suc_idxs = get_all_suc_trajs(env, trajs=all_candidates, goal=target_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Rendering\n",
    "from diffuser.utils.rm2d_render import RandStaticMazeRenderer\n",
    "rd = RandStaticMazeRenderer(plan_env_list)\n",
    "\n",
    "tmp_path = f'{rootdir}/luotest.png'\n",
    "renderer = plan_env_list.model_list[env_id].renderer\n",
    "renderer.fig_dpi = 40 # lower resolution for fast render\n",
    "renderer.up_right = 'default'\n",
    "renderer.config['no_draw_traj_line'] = False\n",
    "renderer.config['no_draw_col_info'] = False\n",
    "\n",
    "# rd.composite( tmp_path, unnm_traj[:10],  np.array([env_id])) # visualize all trajs\n",
    "rd.composite( tmp_path, suc_trajs[:10],  np.array([env_id])) # visualize suc trajs\n",
    "img = imageio.imread(tmp_path)\n",
    "utils.plt_img(img, dpi=200) # 600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
